/*
 *  canny.c
 *  libCV
 *
 *  Created by Gregory Koch on 10/05/11.
 *  Copyright 2011 IFlyBotIV. All rights reserved.
 *
 */

#include "canny.h"

#include "ppm.h"

static void follow_edge(image_t src, image_t dst, int x, int y, char *marked)
{
    int pos_x, pos_y;
    
    pos_x = x-1;
    
    for(pos_y = y-1; pos_y < y + 2; pos_y++)
    {
        if(pos_x >= 0 && pos_x < src.size.width-1 && pos_y >= 0 && pos_y < src.size.height-1)
        {
            if(src.data[src.stride * pos_y + pos_x] > TL && !marked[dst.size.width * pos_y + pos_x])
            {
                marked[dst.size.width * pos_y + pos_x] = 1;
                dst.data[dst.stride * pos_y + pos_x] = 255;
                
                follow_edge(src, dst, pos_x, pos_y, marked);
            }
        }
    }
    
    pos_x = x+1;
    
    for(pos_y = y-1; pos_y < y + 2; pos_y++)
    {
        if(pos_x >= 0 && pos_x < src.size.width-1 && pos_y >= 0 && pos_y < src.size.height-1)
        {
            if(src.data[src.stride * pos_y + pos_x] > TL && !marked[dst.size.width * pos_y + pos_x])
            {
                marked[dst.size.width * pos_y + pos_x] = 1;
                dst.data[dst.stride * pos_y + pos_x] = 255;
                
                follow_edge(src, dst, pos_x, pos_y, marked);
            }
        }
    }
    
    pos_x = x;
    pos_y = y-1;
    
    if(pos_x >= 0 && pos_x < src.size.width-1 && pos_y >= 0 && pos_y < src.size.height-1)
    {
        if(src.data[src.stride * pos_y + pos_x] > TL && !marked[dst.size.width * pos_y + pos_x])
        {
            marked[dst.size.width * pos_y + pos_x] = 1;
            dst.data[dst.stride * pos_y + pos_x] = 255;
            
            follow_edge(src, dst, pos_x, pos_y, marked);
        }
    }
    
    pos_x = x;
    pos_y = y+1;
    
    if(pos_x >= 0 && pos_x < src.size.width-1 && pos_y >= 0 && pos_y < src.size.height-1)
    {
        if(src.data[src.stride * pos_y + pos_x] > TL && !marked[dst.size.width * pos_y + pos_x])
        {
            marked[dst.size.width * pos_y + pos_x] = 1;
            dst.data[dst.stride * pos_y + pos_x] = 255;
            
            follow_edge(src, dst, pos_x, pos_y, marked);
        }
    }

    
}

void canny(image_t src, image_t dst)
{
    int m_x, m_y, x, y, pos_src, pos_dst, pos_tmp, val_x, val_y;
    float *theta;
    char *marked;
    
    image_t tmp;
    
    int gaussian[] = {2, 4, 5, 4, 2, 4, 9, 12, 9, 4, 5, 12, 15, 12, 5, 4, 9, 12, 9, 4, 2, 4, 5, 4, 2};
    
    int sobel_x[] = {1, 2, 1, 0, 0, 0, -1, -2, -1};
    int sobel_y[] = {1, 0, -1, 2, 0, -2, 1, 0, -1};
    
    if(src.type != IMG_GREYSCALE8 && dst.type != IMG_GREYSCALE8)
	{
		perror("libCV : canny works only with greyscale images");
		exit(EXIT_FAILURE);
	}
    
    if(src.size.width != dst.size.width || src.size.height != dst.size.height)
	{
		perror("libCV : canny : src and dst should have the same size");
		exit(EXIT_FAILURE);
	}
    
    tmp = init_image(IMG_GREYSCALE8, init_size(src.size.width, src.size.height));
    
    theta = malloc(sizeof(int) * src.size.width * src.size.height);
    if(theta == NULL)
    {
        perror("libCV : malloc error");
        exit(EXIT_FAILURE);
    }

    marked = malloc(sizeof(char) * src.size.width * src.size.height);
    if(theta == NULL)
    {
        perror("libCV : malloc error");
        exit(EXIT_FAILURE);
    }
    
    memset(marked, 0, sizeof(char) * src.size.width * src.size.height);
    
    /*
     * First we apply the 5*5 gaussian mask
     */
    for(y=src.size.height-1; y--; )
    {
        for(x=src.size.width-1; x--; )
        {
            val_x = 0;
            
            for(m_y=0; m_y<5; m_y++)
            {
                for(m_x=0; m_x<5; m_x++)
                {
                    pos_src = src.stride * (y - 2 + m_y) + x - 2 + m_x;
                    if(pos_src >= 0 && pos_src < (src.size.height * src.stride))
                    {
                        val_x += src.data[pos_src] * gaussian[5*m_y + m_x];
                    }
                }
            }
            
            val_x /= 159;
            
            val_x = val_x > 255 ? 255 : val_x; 
            
            pos_tmp = tmp.stride * y + x;
            
            tmp.data[pos_tmp] = val_x;
        }
    }
    
    /*
     * Sobel 3*3
     */
    for(y=tmp.size.height-1; y--; )
    {
        for(x=tmp.size.width-1; x--; )
        {
            val_x = 0;
            val_y = 0;
            
            for(m_y=0; m_y<3; m_y++)
            {
                for(m_x=0; m_x<3; m_x++)
                {
                    pos_tmp = tmp.stride * (y - 1 + m_y) + x - 1 + m_x;
                    if(pos_tmp >= 0 && pos_tmp < (tmp.size.height * tmp.stride))
                    {
                        val_x += tmp.data[pos_tmp] * sobel_x[3*m_y + m_x];
                        val_y += tmp.data[pos_tmp] * sobel_y[3*m_y + m_x];
                    }
                }
            }
            
            val_x = val_x >= 0 ? val_x : -val_x;
            val_y = val_y >= 0 ? val_y : -val_y; 
            
            pos_dst = dst.stride * y + x;
            
            dst.data[pos_dst] = val_x + val_y;
            
            if(val_x != 0)
                theta[y * tmp.size.width + x] =  atan(val_y/val_x);
            else
            {
                if(val_y)
                    theta[y * tmp.size.width + x] =  90;
                else
                    theta[y * tmp.size.width + x] =  0;
            }
        }
    }

    /*
     * Non Maxima Suppression
     */
    for(y=dst.size.height--; y--; )
    {
        for(x=dst.size.width-1; x--; )
        {
            pos_dst = dst.stride * y + x;
            pos_tmp = tmp.stride * y + x;
            
            if(theta[y * dst.size.width + x] == -1)
            {
                tmp.data[pos_tmp] = 0;
            }
            else if((theta[y * dst.size.width + x] >= 0 && theta[y * dst.size.width + x] < 22.5) 
               || (theta[y * dst.size.width + x] >= 157.5 && theta[y * dst.size.width + x] < 202.5)
               || (theta[y * dst.size.width + x] >= 337.5 && theta[y * dst.size.width + x] <= 360))
            {
                if(y == 0)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else if(y == dst.size.height - 1)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x] 
                       && dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
            }
            else if((theta[y * dst.size.width + x] >= 22.5 && theta[y * dst.size.width + x] < 67.5) 
                    || (theta[y * dst.size.width + x] >= 202.5 && theta[y * dst.size.width + x] < 247.5))
            {
                if(y == 0 && x == dst.size.width-1)
                {
                    tmp.data[pos_tmp] = 0;
                }
                else if(y == dst.size.height-1 && x == 0)
                {
                    tmp.data[pos_tmp] = 0;
                }
                else if(y == 0 || x == 0)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x+1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else if(y == dst.size.height - 1 || x == dst.size.width-1)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x-1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x+1] 
                       && dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x-1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
            }
            else if((theta[y * dst.size.width + x] >= 67.5 && theta[y * dst.size.width + x] < 112.5 )
                    || (theta[y * dst.size.width + x] >= 247.5 && theta[y * dst.size.width + x] < 292.5))
            {
                if(x == 0)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * y + x + 1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else if(x == dst.size.width - 1)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * y + x - 1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * y + x - 1] 
                       && dst.data[pos_dst] > dst.data[dst.stride * y + x + 1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                } 
            }
            else if((theta[y * dst.size.width + x] >= 112.5 && theta[y * dst.size.width + x] < 157.5) 
                    || (theta[y * dst.size.width + x] >= 292.5 && theta[y * dst.size.width + x] < 337.5))
            {
                if(y == 0 && x == 0)
                {
                    tmp.data[pos_tmp] = 0;
                }
                else if(y == dst.size.height-1 && x == dst.size.width-1)
                {
                    tmp.data[pos_tmp] = 0;
                }
                else if(y == 0 || x == dst.size.width-1)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x-1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else if(y == dst.size.height - 1 || x == 0)
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x+1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
                else
                {
                    if(dst.data[pos_dst] > dst.data[dst.stride * (y-1) + x-1] 
                       && dst.data[pos_dst] > dst.data[dst.stride * (y+1) + x+1])
                    {
                        tmp.data[pos_tmp] = dst.data[pos_dst];
                    }
                    else
                    {
                        tmp.data[pos_tmp] = 0;                        
                    }
                }
            }
                        
        }
    }

    /*
     * Hysteresis Thresholding
     */
    for(y=tmp.size.height-1; y--; )
    {
        for(x=tmp.size.width-1; x--; )
        {
            pos_tmp = tmp.stride * y + x;
            
            if(!marked[dst.size.width * y + x])
            {
                if(tmp.data[pos_tmp] > TH)
                {
                    marked[dst.size.width * y + x] = 1;
                    dst.data[dst.stride * y + x] = 255;
                    
                    follow_edge(tmp, dst, x, y, marked);
                }
                else
                {
                    dst.data[dst.stride * y + x] = 0;
                }
            }
            
        }
    }
    
    free(marked);
    free(theta);
    free_image(tmp);
}
